#!/usr/bin/env python3
"""Describe `data/国家级生态乡镇名单.parquet` and check uniqueness by (县,乡镇,年份).

Outputs a human-readable summary to stdout and saves JSON/MD reports to `data/outputs/`.
"""
from __future__ import annotations

import json
from pathlib import Path
from datetime import datetime
import math

try:
    import polars as pl
except Exception:
    pl = None


def round2(x):
    try:
        return None if x is None or (isinstance(x, float) and math.isnan(x)) else round(float(x), 2)
    except Exception:
        return x


def main():
    if pl is None:
        print('polars not available in this environment; aborting')
        return

    src = Path('data/国家级生态乡镇名单.parquet')
    out_dir = Path('data/outputs')
    out_dir.mkdir(parents=True, exist_ok=True)

    df = pl.read_parquet(src)
    total_rows = df.height

    # Uniqueness check by (县,乡镇,年份)
    grp = df.group_by(['县', '乡镇', '年份'], maintain_order=True).agg(pl.count().alias('cnt'))
    duplicates = grp.filter(pl.col('cnt') > 1).sort('cnt', reverse=True)

    # Column-level stats
    cols = df.columns
    col_stats = {}
    for c in cols:
        s = df[c]
        non_nulls = s.drop_nulls().height
        nulls = total_rows - non_nulls
        distincts = s.n_unique()
        dtype = s.dtype
        info = {
            'non_nulls': int(non_nulls),
            'nulls': int(nulls),
            'distincts': int(distincts),
            'dtype': str(dtype),
        }

        if dtype == pl.Utf8:
            # top 5 values
            vc = df.groupby(c).agg(pl.count().alias('cnt')).sort('cnt', reverse=True).head(5)
            info['top5'] = [{ 'value': r[c], 'count': int(r['cnt']) } for r in vc.to_dicts()]
        else:
            try:
                info.update({
                    'min': round2(df[c].min()),
                    'median': round2(df[c].quantile(0.5)),
                    'mean': round2(df[c].mean()),
                    'std': round2(df[c].std()),
                })
            except Exception:
                pass

        col_stats[c] = info

    # Prepare outputs
    ts = datetime.now().strftime('%Y%m%dT%H%M%S')
    summary = {
        'file': str(src),
        'total_rows': int(total_rows),
        'duplicates_count': int(duplicates.height),
        'duplicates_preview': duplicates.head(20).to_dicts(),
        'columns': col_stats,
    }

    json_path = out_dir / f'国家级生态乡镇名单_describe_{ts}.json'
    md_path = out_dir / f'国家级生态乡镇名单_describe_{ts}.md'

    with json_path.open('w', encoding='utf8') as f:
        json.dump(summary, f, ensure_ascii=False, indent=2)

    # write a simple markdown summary
    with md_path.open('w', encoding='utf8') as f:
        f.write(f"# 描述统计报告 — 国家级生态乡镇名单\n\n")
        f.write(f"生成时间: {ts}\n\n")
        f.write(f"- 文件: `{src}`\n")
        f.write(f"- 总行数: {total_rows}\n")
        f.write(f"- 按 (县,乡镇,年份) 重复组数: {duplicates.height}\n\n")

        if duplicates.height > 0:
            f.write('## 重复样例 (最多20)\n')
            f.write('| 县 | 乡镇 | 年份 | cnt |\n')
            f.write('|---|---:|---:|---:|\n')
            for r in duplicates.head(20).to_dicts():
                f.write(f"| {r['县']} | {r['乡镇']} | {r['年份']} | {r['cnt']} |\n")
            f.write('\n')

        f.write('## 列级统计\n')
        for c, info in col_stats.items():
            f.write(f"### {c}\n")
            f.write(f"- dtype: {info.get('dtype')}\n")
            f.write(f"- non_nulls: {info.get('non_nulls')}\n")
            f.write(f"- nulls: {info.get('nulls')}\n")
            f.write(f"- distincts: {info.get('distincts')}\n")
            if 'top5' in info:
                f.write('- top5:\n')
                f.write('| value | count |\n')
                f.write('|---|---:|\n')
                for it in info['top5']:
                    f.write(f"| {it['value']} | {it['count']} |\n")
            else:
                if 'mean' in info:
                    f.write(f"- min: {info.get('min')}\n")
                    f.write(f"- median: {info.get('median')}\n")
                    f.write(f"- mean: {info.get('mean')}\n")
                    f.write(f"- std: {info.get('std')}\n")
            f.write('\n')

    # Print concise summary to stdout
    print(f"File: {src}")
    print(f"Total rows: {total_rows}")
    print(f"Duplicate groups (县,乡镇,年份) count: {duplicates.height}")
    if duplicates.height > 0:
        print('\nTop duplicate groups (最多20):')
        print(duplicates.head(20).to_pandas())

    print('\nPer-column summary:')
    for c, info in col_stats.items():
        print(f"- {c}: dtype={info.get('dtype')}, non_nulls={info.get('non_nulls')}, nulls={info.get('nulls')}, distincts={info.get('distincts')}")
        if 'top5' in info:
            for it in info['top5']:
                print(f"    top: {it['value']} ({it['count']})")

    print('\nSaved JSON:', json_path)
    print('Saved MD:', md_path)


if __name__ == '__main__':
    main()
